use std::iter::Peekable;
use std::str::Chars;

// ----------------------
//  AST 数据结构
// ----------------------
#[derive(Debug, PartialEq)]
enum Expr {
   Number(f64),
   Prefix(Op, Box<Expr>),
   Infix(Op, Box<Expr>, Box<Expr>),
}

#[derive(Debug, PartialEq, Clone, Copy)]
enum Op {
   Add,
   Subtract,
   Multiply,
   Divide,
   Power,
   Minus, // 负号
}

// ----------------------
//  词法分析器 (Lexer)
// ----------------------
#[derive(Debug, Clone, Copy, PartialEq)]
enum Token {
   Number(f64),
   Plus,
   Minus,
   Star,
   Slash,
   Caret,
   LeftParen,
   RightParen,
   EOF,
}

struct Lexer<'a> {
   chars: Peekable<Chars<'a>>,
}

impl<'a> Lexer<'a> {
   fn new(input: &'a str) -> Self {
      Lexer {
         chars: input.chars().peekable(),
      }
   }

   fn next_token(&mut self) -> Token {
      self.skip_whitespace();

      match self.chars.next() {
         Some('+') => Token::Plus,
         Some('-') => Token::Minus,
         Some('*') => Token::Star,
         Some('/') => Token::Slash,
         Some('^') => Token::Caret,
         Some('(') => Token::LeftParen,
         Some(')') => Token::RightParen,
         Some(c) if c.is_ascii_digit() || c == '.' => {
            let mut num = String::new();
            num.push(c);
            while let Some(&next) = self.chars.peek() {
               if next.is_ascii_digit() || next == '.' {
                  num.push(self.chars.next().unwrap());
               } else {
                  break;
               }
            }
            Token::Number(num.parse().unwrap())
         }
         None => Token::EOF,
         _ => panic!("Unexpected character"),
      }
   }

   fn skip_whitespace(&mut self) {
      while let Some(&c) = self.chars.peek() {
         if c.is_whitespace() {
            self.chars.next();
         } else {
            break;
         }
      }
   }
}

// ----------------------
//  Pratt 解析器 (Parser)
// ----------------------
struct Parser<'a> {
   lexer: Lexer<'a>,
   current: Token,
}

impl<'a> Parser<'a> {
   fn new(mut lexer: Lexer<'a>) -> Self {
      let current = lexer.next_token();
      Parser { lexer, current }
   }

   fn parse(&mut self) -> Expr {
      self.parse_expression(0)
   }

   fn parse_expression(&mut self, min_bp: u8) -> Expr {
      let mut left = self.parse_prefix();

      loop {
         let op = match self.current {
            Token::Plus => Op::Add,
            Token::Minus => Op::Subtract,
            Token::Star => Op::Multiply,
            Token::Slash => Op::Divide,
            Token::Caret => Op::Power,
            _ => break,
         };

         let (left_bp, right_bp) = self.get_binding_power(op);
         if left_bp < min_bp {
            break;
         }

         self.advance();
         let right = self.parse_expression(right_bp);
         left = Expr::Infix(op, Box::new(left), Box::new(right));
      }

      left
   }

   fn parse_prefix(&mut self) -> Expr {
      let token = self.current;
      self.advance();

      match token {
         Token::Number(n) => Expr::Number(n),
         Token::Minus => {
            let expr = self.parse_expression(30); // 负号优先级最高
            Expr::Prefix(Op::Minus, Box::new(expr))
         }
         Token::LeftParen => {
            let expr = self.parse_expression(0);
            self.consume(Token::RightParen);
            expr
         }
         _ => panic!("Unexpected token: {:?}", token),
      }
   }

   fn get_binding_power(&self, op: Op) -> (u8, u8) {
      match op {
         Op::Add | Op::Subtract => (10, 20),
         Op::Multiply | Op::Divide => (20, 30),
         Op::Power => (40, 39), // 右结合
         Op::Minus => (30, 0),  // 前缀运算符
      }
   }

   fn advance(&mut self) {
      self.current = self.lexer.next_token();
   }

   fn consume(&mut self, expected: Token) {
      if self.current == expected {
         self.advance();
      } else {
         panic!("Expected {:?}, found {:?}", expected, self.current);
      }
   }
}

// ----------------------
//  使用示例
// ----------------------
fn main() {
   let input = "3 + 4 * 2 / (1 - 5)^2";
   let lexer = Lexer::new(input);
   let mut parser = Parser::new(lexer);
   let ast = parser.parse();

   println!("{:#?}", ast);
}
